/**
 * (C) Copyright MarketLive. 2007. All rights reserved.
 * MarketLive is a trademark of MarketLive, Inc.
 * Warning: This computer program is protected by copyright law and international treaties.
 * Unauthorized reproduction or distribution of this program, or any portion of it, may result
 * in severe civil and criminal penalties, and will be prosecuted to the maximum extent
 * possible under the law.
 */
package com.marketlive.system.keymanagement;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.lang.management.ManagementFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.KeyStore.SecretKeyEntry;
import java.security.cert.CertificateException;
import java.security.spec.InvalidKeySpecException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;

/**
 * This class is responsible for creating the JCEKS keystore,
 * generating and managing all symmetric keys in this keystore.
 */
public class KeystoreManager {

    public static final String KS_TYPE = "JCEKS";
    public static final String KS_FILENAME = "marketlive.keystore";
    public static final int KEY_SIZE = 256;
    public static final String HASH_TYPE = "SHA-256";
    public static final String AES_KEY_TYPE = "AES";
    public static final String RSA_KEY_TYPE = "RSA";
    public static final String PBE_ALGORITHM = "PBEWithSHA1AndDESede";
    public Connection con = null;
    
    /**
     * Constructor
     */
    public KeystoreManager() {
        // EMPTY
    }
    
    /**
     * Creates the JCEKS KeyStore
     * 
     * @param   location    Directory where the keystore is to be
     * created; only directory name ending with a trailing "/" is 
     * required - keystore name is not to be provided
     * @param   password that protects the keystore
     * 
     * @return boolean true if the method succeeded, or false otherwise
     * 
     * @exception InvalidLocationException if location is zero-length
     * @exception InvalidPasswordException if password is zero-length
     * @exception DuplicateKeystoreException if keystore already exists
     * @exception OtherKeyMgtException for all other exceptions caught
     */
    public boolean createKeystore(String location, char[] password)
            throws InvalidPasswordException, InvalidLocationException, 
            DuplicateKeystoreException, OtherKeyMgtException {        
        
        // Is location zero-length?
        if (location.length() == 0)
            throw new InvalidLocationException("Keystore location cannot be zero-length");
        else if (location.charAt(location.length() - 1) != '/')
            throw new InvalidLocationException("Keystore location must be a directory and end with a /");
        
        // Is password zero-length?
        if (password.length == 0)
            throw new InvalidPasswordException("Keystore password cannot be zero-length");
        
        // Does keystore exist?
        File f = new File(location + KS_FILENAME);
        if (f.exists())
            throw new DuplicateKeystoreException(
                    "Keystore already exists at this location: " + f.getAbsolutePath());
        
        // Create keystore 
        KeyStore ks;
        
        // Write out keystore 
        try {
            ks = KeyStore.getInstance(KS_TYPE);
            ks.load(null, password);
            ks.store(new FileOutputStream(f),password);
        } catch (KeyStoreException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A KeyStoreException was thrown; check logs");
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A NoSuchAlgorithmException was thrown; check logs");
        } catch (CertificateException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A CertificateException was thrown; check logs");
        } catch (FileNotFoundException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A FileNotFoundException was thrown; check logs");
        } catch (IOException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("An IOException was thrown; check logs");
        }
        
        return true;
    }
    
    /**
     * Generates a cryptographic key; currently only symmetric keys
     * are needed (is this true?), so only AES-256 symmetric keys are 
     * generated.
     * 
     * @param String location Directory where the keystore is to be 
     * created; only directory name ending with a trailing "/" is 
     * required - keystore name is not to be provided
     * @param char[] password that protects the keystore 
     * @param String alias The identifier for the key in the keystore
     * @param String keytype The type of key to be generated: DES, AES, etc.
     * @param short size The size of the key: 128, 192, 256, etc.
     * 
     * @return boolean true if the method succeeded, or false otherwise
     * 
     * @exception InvalidLocationException if location is zero-length
     * @exception InvalidPasswordException if password is zero-length
     * @exception InvalidAliasException if the alias already exists
     * @exception FileNotFoundException if the keystore cannot be found
     * @exception OtherKeyMgtException for all other exceptions
     */
    public boolean generateKey(String location, char[] password, 
            String alias, String keytype, short size) 
            throws InvalidPasswordException, InvalidLocationException, 
            InvalidAliasException, FileNotFoundException, OtherKeyMgtException {
                
        // Is location zero-length?
        if (location.length() == 0)
            throw new InvalidLocationException("Keystore location cannot be zero-length");
        
        // Is password zero-length?
        if (password.length == 0)
            throw new InvalidPasswordException("Keystore password cannot be zero-length");
        
        // Is alias zero-length?
        if (alias.length() == 0)
            throw new InvalidAliasException("Key alias cannot be zero-length");
        
        // Is keytype zero-length?
        if (keytype.length() == 0)
            ;
            // Doesn't matter because we only generate AES keys
            // Need to later add a check for RSA key-type
        
        // Is keytype zero-length?
        if (size == 0)
            ;
            // Doesn't matter because we only generate AES-256 keys
            // Need to later add a check for RSA key-size
        
        // Does keystore exist at location?
        File ksfile = new File(location + KS_FILENAME);
        if (!ksfile.exists())
            throw new FileNotFoundException("Keystore not found at location: " + location);
        
        // Create keystore
        KeyStore ks;
        
        try {
            ks = KeyStore.getInstance(KS_TYPE);
        
            // Read keystore
            FileInputStream ksfis = new FileInputStream(ksfile);
        
            // Load keystore
            ks.load(ksfis, password);
            
            // Close input stream
            ksfis.close();

            // Is alias already defined in keystore?
            if (ks.containsAlias(alias))
                throw new InvalidAliasException(
                        "A key with this alias was found in the keystore: " + alias);
        
            // Generate the key
            KeyGenerator keyGen = KeyGenerator.getInstance(AES_KEY_TYPE);
            
            // Initialize generator
            keyGen.init(KEY_SIZE, new SecureRandom());
            
            // Generate key
            SecretKey aesKey = keyGen.generateKey();
            
            // Create entry in keystore for new key
            SecretKeyEntry skEntry = new SecretKeyEntry(aesKey);
            
            // Save key in entry
            ks.setEntry(alias, skEntry, new KeyStore.PasswordProtection(password));
            
            /**
             * Write out keystore with temporary filename
             * The getName() from ManagementFactory returns the PID
             */ 
            File tmpFile = new File(location + KS_FILENAME + 
                    ManagementFactory.getRuntimeMXBean().getName());
            ks.store(new FileOutputStream(tmpFile), password);
            
            // Delete the .old keystore if it exists  
            File oldFile = new File (location + KS_FILENAME + ".old");
            if (oldFile.exists())
                oldFile.delete();
            
            // Rename existing keystore file to .old
            ksfile.renameTo(oldFile);
            tmpFile.renameTo(ksfile);    
                
        } catch (KeyStoreException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A KeyStoreException was thrown; check logs");
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A NoSuchAlgorithmException was thrown; check logs");
        } catch (CertificateException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A CertificateException was thrown; check logs");
        } catch (IOException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A IOException was thrown; check logs");        
        } 
        
        return true;
    }
 
    /**
     * Generates a PBE cryptographic key; currently only a
     * PBEWithSHA1AndDESede key is supported
     * 
     * @param String keycode the code that serves as password 
     * @return SecretKey if the method succeeded
     * 
     * @exception NoSuchAlgorithmException if the PBE algorithm is unavailable
     * @exception InvalidKeySpecException if the keyspec is invalid
     * @exception UnsupportedEncodingException if the UTF-8 encoding is unsupported
     * 
     */
    public SecretKey generatePBEKey(String keycode) 
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        
        // Get keycode into a char array
        char[] kcchars = new char[keycode.length()];
        keycode.getChars(0, keycode.length(), kcchars, 0);
        
//        System.out.println("char password: " + new String(kcchars));
        
        // Create a PBEKeySpec & SecretKeyFactory
        PBEKeySpec pbeks = new PBEKeySpec(kcchars);
        SecretKeyFactory skf = SecretKeyFactory.getInstance(PBE_ALGORITHM);
        
        // Generate and return key
        return skf.generateSecret(pbeks);
    }
    
    /**
     * Generates an RSA cryptographic key-pair - a public and private
     * key - for use in encrypting PAN and Customer passwords;  
     * 
     * @param short keysize The size of the key: 1024, 2048, 3072 or 4096
     * @return KeyPair containing the RSA cryptographic keys
     * 
     * @exception NoSuchAlgorithmException if the PBE algorithm is unavailable
     * @exception OtherKeyMgtException for all other exceptions
     */
    public KeyPair generateRSAKey(short keysize) 
            throws OtherKeyMgtException {
        
        // Is keysize zero-length? Use 1024 as default
        if (keysize == 0)
            keysize = 1024;
        
        switch (keysize) {
            case 1024: break;
            case 2048: break;
            case 3072: break;
            case 4096: break;
            default:
                throw new OtherKeyMgtException("RSA Keysize is not 1024, 2048, 3072 or 4096: " + keysize);
        }
        
        try {
        
            // Create the keypair generator
            KeyPairGenerator kpg = KeyPairGenerator.getInstance(RSA_KEY_TYPE);
            System.out.println("Using JCE Provider: " + kpg.getProvider().getName());
            
            // Initialize it with the appropriate parameters
            kpg.initialize(keysize, new SecureRandom());
            
            // Generate the keys
            System.out.println("Generated RSA key-pair successfully");
            return kpg.generateKeyPair();

        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
            throw new OtherKeyMgtException("A NoSuchAlgorithmException was thrown; check logs");
        }
    }
    
    /**
     * Checks if a key is in the keystore
     * 
     * @param String location Directory where the keystore is to be 
     * created; only directory name ending with a trailing "/" is 
     * required - keystore name is not to be provided
     * @param String alias The identifier for the key in the keystore
     * @return boolean true if the is in the keystore, or false, if not
     * @exception Exception if there is some other problem from the JCE
     */
    public boolean isKeyinKeystore(String location, String alias) 
                throws Exception {
        return false;
    }
       
    /**
     * Lists keys in the keystore
     * 
     * @param String location Directory where the keystore is to be 
     * created; only directory name ending with a trailing "/" is 
     * required - keystore name is not to be provided
     * @return String[] Array of key-aliases from the keystore
     * @exception Exception if there is some other problem from the JCE
     */
    public String[] listKeys(String location) throws Exception {
        return null;
    }
    
    /**
     * Gets a key from the JCEKS keystore
     * 
     * @param String Directory location where the keystore exists 
     * @param char[] password that protects the keystore
     * @param String alias The identifier for the key in the keystore
     * 
     * @return byte[] The key.  The assumption is that the caller knows
     * what type of key is being returned back, based on the alias
     * @exception Exception if the keystore password is incorrect, or if 
     * there is some other problem from the JCE
     */
    public byte[] getKey(String location, char[] password,
            String alias) throws Exception {
        return null;
    }
    
    /**
     * Overloaded method to get a key from the SYSTEM_KEYSTORE in RDBMS
     * 
     * @param String keycode that identifies the required key 
     * @param DBPropertiesBean for the JDBC paramethers
     * @param String keytype - either "public" or "private"
     * 
     * @return String containing the Base64-encoded RSA Key
     */
    public String getKey(String keycode, DBPropertiesBean dbpb, String keytype) {
        
        // Local variable
        PreparedStatement stmt = null;
        
        try {
            // Setup connection
            if (con == null || con.isClosed()) {
                con = Utilities.getDBConnection(dbpb);
            }

            // Setup SQL to retrieve key
            if ("private".equals(keytype)) {
                stmt = con.prepareStatement(
                    "SELECT PRIVATE_KEY FROM SYSTEM_KEYSTORE WHERE KEY_CODE = ?");
            } else { // PublicKey
                stmt = con.prepareStatement(
                    "SELECT PUBLIC_KEY FROM SYSTEM_KEYSTORE WHERE KEY_CODE = ?");
            }

            // Provide search parameter - should return only 1 record 
            stmt.setString(1, keycode);     // Keycode
            
            // Execute and get base64-encoded, encrypted PrivateKey
            ResultSet rs = stmt.executeQuery();
            rs.next();
            if ("private".equals(keytype)) {
                String b64pvk = rs.getString(1);
                if (b64pvk != null) {
                    return b64pvk;
                }
            } else {    // PublicKey
                String b64pbk = rs.getString(1);
                if (b64pbk != null) {
                    return b64pbk;
                }
            }
            
        } catch (ClassNotFoundException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (SQLException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (NoSuchAlgorithmException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (UnsupportedEncodingException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        }
        
        // Failed 
        return null;
    }
    
    /**
     * Deletes a key from the keystore
     * 
     * @param String location Directory where the keystore is to be 
     * created; only directory name ending with a trailing "/" is 
     * required - keystore name is not to be provided
     * @param char[] password that protects the keystore
     * @param String alias The identifier for the key in the keystore
     * 
     * @return boolean true if the method succeeded, or false otherwise
     * @exception Exception if the keystore password is incorrect, or if 
     * there is some other problem from the JCE
     */
    public boolean deleteKey(String location, char[] password,
            String alias) throws Exception {
        return false;
    }
    
    /**
     * Changes the keystore's password.  
     * 
     * @param String location Directory where the keystore is to be 
     * created; only directory name ending with a trailing "/" is 
     * required - keystore name is not to be provided
     * @param char[] the old password that protects the keystore
     * @param char[] the new password to protect the keystore
     * @return boolean true if the is in the keystore, or false, if not
     * @exception Exception if the keystore password is incorrect, or if
     * there is some other problem from the JCE
     */
    public boolean changePassword(String location, char[] oldpassword, 
            char[] newpassword) throws Exception {
        return false;
    }
    
    /**
     * Uses JDBC to connect to the database, and save a
     * key code and public/private key pair.
     * 
     * @param   keycode
     * @param   publickey
     * @param   privatekey
     * @param   dbpb
     *
     * @return boolean true if the method succeeded, or false otherwise
     */
    protected boolean saveKeysToDB(String keycode, String publickey, 
            String privatekey, DBPropertiesBean dbpb) {
        
        try {
            // Setup connection
            if (con == null || con.isClosed()) {
                con = Utilities.getDBConnection(dbpb);
            }
            
            // Create SQL statement
            PreparedStatement stmt = con.prepareStatement(
                    "INSERT INTO SYSTEM_KEYSTORE (KEY_CODE, VERSION, PRIVATE_KEY," +
                    "PUBLIC_KEY, DATE_CREATED, ACTIVE) VALUES(?, ?, ?, ?, ?, ?)");

            // Update values in SQL that need to be inserted
            stmt.setString(1, keycode);     // Keycode
            stmt.setInt(2, 0);              // Version
            stmt.setString(3, privatekey);  // Private key
            stmt.setString(4, publickey);   // Public key
            // Current date and time
            stmt.setTimestamp(5, new java.sql.Timestamp(new java.util.Date().getTime()));
            stmt.setInt(6, 1);              // Active
            
            // Save to DB; check for success
            if (stmt.executeUpdate() == 1) {
                    System.out.println("New RSA keys were stored to DB successfully");
                    return true;
            }
            
        } catch (ClassNotFoundException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (SQLException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (NoSuchAlgorithmException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (UnsupportedEncodingException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        }
        
        // Failed 
        System.out.println("Keys were NOT stored to DB; check logs for errors");
        return false;
    }
    
    /**
     * Sets a currently inactive key-pair to active status
     * 
     * @param String keycode that identifies the required key 
     * @param DBPropertiesBean for the JDBC paramethers
     * 
     * @return boolean indicating success (true) or failure (false)
     */
    public boolean setKeyActive(String keycode, DBPropertiesBean dbpb) {
        
        try {
            // Setup connection
            if (con == null || con.isClosed()) {
                con = Utilities.getDBConnection(dbpb);
            }
            
            // Setup SQL to retrieve key to check status first
            PreparedStatement stmt = con.prepareStatement(
                    "SELECT ACTIVE FROM SYSTEM_KEYSTORE WHERE KEY_CODE = ?");

            // Provide search parameter - should return only 1 record 
            stmt.setString(1, keycode);     // Keycode
            
            // Execute and get base64-encoded, encrypted PrivateKey
            ResultSet rs = stmt.executeQuery();
            rs.next();
            int status = rs.getInt(1);
            if (status != 0) { 
                System.err.println("\nERROR: Keycode " + keycode + " is not inactive; cannot change status to active");
                return false;
            }
            
            // Otherwise continue to update status
            stmt = con.prepareStatement(
                    "UPDATE SYSTEM_KEYSTORE SET ACTIVE = 1 WHERE KEY_CODE = ?");

            // Provide search parameter  
            stmt.setString(1, keycode);     // Keycode
            
            // Update DB; check for success
            if (stmt.executeUpdate() == 1) {
                    System.out.println("Keycode " + keycode + " was set to Active successfully");
                    return true;
            }
            
        } catch (ClassNotFoundException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (SQLException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (NoSuchAlgorithmException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (UnsupportedEncodingException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        }
        
        // Failed 
        return false;
    }
    
    
    /**
     * Sets a currently active key-pair to inactive status
     * 
     * @param String keycode that identifies the required key 
     * @param DBPropertiesBean for the JDBC paramethers
     * 
     * @return boolean indicating success (true) or failure (false)
     */
    public boolean setKeyInactive(String keycode, DBPropertiesBean dbpb) {
        
        try {
            // Setup connection
            if (con == null || con.isClosed()) {
                con = Utilities.getDBConnection(dbpb);
            }
            
            // Setup SQL to retrieve key to check status first
            PreparedStatement stmt = con.prepareStatement(
                    "SELECT ACTIVE FROM SYSTEM_KEYSTORE WHERE KEY_CODE = ?");

            // Provide search parameter - should return only 1 record 
            stmt.setString(1, keycode);     // Keycode
            
            // Execute and get base64-encoded, encrypted PrivateKey
            ResultSet rs = stmt.executeQuery();
            rs.next();
            int status = rs.getInt(1);
            
            // If status is already inactive, do nothing
            if (status == 0) { 
                System.out.println("\nKeycode " + keycode + " is already inactive; status was not changed");
                return true;
            }
            
            // Setup SQL to retrieve key
            stmt = con.prepareStatement(
                    "UPDATE SYSTEM_KEYSTORE SET ACTIVE = 0 WHERE KEY_CODE = ?");

            // Provide search parameter  
            stmt.setString(1, keycode);     // Keycode
            
            // Update DB; check for success
            if (stmt.executeUpdate() == 1) {
                    System.out.println("Keycode " + keycode + " was set to Inactive successfully");
                    return true;
            }
            
        } catch (ClassNotFoundException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (SQLException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (NoSuchAlgorithmException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (UnsupportedEncodingException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        }
        
        // Failed 
        return false;
    }
    
    /**
     * Sets all crypto keys in database to inactive status; this is
     * necessary for new installations that may not have cleaned out
     * the test keys from database
     *  
     * @param DBPropertiesBean for the JDBC paramethers
     * 
     * @return boolean indicating success (true) or failure (false)
     */
    public boolean setAllKeysInactive(DBPropertiesBean dbpb) {
        
        try {
            // Setup connection
            if (con == null || con.isClosed()) {
                con = Utilities.getDBConnection(dbpb);
            }
            
            // Setup SQL to update status
            PreparedStatement stmt = con.prepareStatement(
                    "UPDATE SYSTEM_KEYSTORE SET ACTIVE = 0 WHERE ACTIVE = 1");
            
            // Update DB; check for success
            int n = 0;
            if ((n = stmt.executeUpdate()) > 0) {
                    System.out.println("Updated " + n + " old key(s) to Inactive status successfully");
                    return true;
            } else if (n == 0) {
                System.out.println("There were no Active keys to update to Inactive status");
                return true;
            }
            
        } catch (ClassNotFoundException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (SQLException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (NoSuchAlgorithmException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        } catch (UnsupportedEncodingException e) {
            System.err.println("\nERROR: " + e.getLocalizedMessage() + '\n');
            e.printStackTrace();
        }
        
        // Failed 
        return false;
    }
}
